package com.niodata.dp.task.az;

import com.google.common.base.Strings;
import com.niodata.dp.util.HttpCaller;
import com.niodata.dp.util.JsonUtil;
import java.io.File;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import org.apache.commons.lang3.StringEscapeUtils;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.mime.HttpMultipartMode;
import org.apache.http.entity.mime.MultipartEntityBuilder;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.util.EntityUtils;

public class AzkabanApi {

  private static final String URL_LOGIN = "/?action=login";
  private static final String URL_CREATE_PROJECT = "/manager?action=create";
  private static final String URL_DELETE_PROJECT = "/manager?delete=true";
  private static final String URL_UPLOAD_PROJECT = "/manager?ajax=upload";
  private static final String URL_GET_PROJECT_FLOWS = "/manager?ajax=fetchprojectflows";
  private static final String URL_GET_JOBS_OF_FLOW = "/manager?ajax=fetchflowgraph";
  private static final String URL_SET_JOB_PARAMS = "/manager?ajax=setJobOverrideProperty";
  private static final String URL_GET_JOB_PARAMS = "/manager?ajax=getJobProperty";
  private static final String URL_EXECUTE_FLOW = "/executor?ajax=executeFlow";
  private static final String URL_CALCEL_FLOW_EXECUTION = "/executor?ajax=cancelFlow";
  private static final String URL_GET_EXEC_JOB_LOGS = "/executor?ajax=fetchExecJobLogs";
  private static final String URL_GET_FLOW_EXECUTIONS = "/manager?ajax=fetchFlowExecutions";
  private static final String URL_GET_FLOW_RUNNING_EXECUTIONS = "/executor?ajax=getRunning";
  private static final String URL_SCHEDULE_FLOW = "/schedule?ajax=scheduleCronFlow";
  private static final String URL_GET_SCHEDULE = "/schedule?ajax=fetchSchedule";
  private static final String URL_REMOVE_SCHEDULE = "/schedule?action=removeSched";
  private static final String URL_GET_FLOW_EXEC_DETAIL = "/executor?ajax=fetchexecflow";
  private static final String URL_GET_ACTIVE_EXECUTORS = "/executor?ajax=fetchActiveExecutors";
  private static final String URL_GET_RECENTLY_FINISHED = "/executor?ajax=recentlyFinished";
  private static final String URL_GET_RUNNING_FLOWS = "/executor?ajax=runningFlows";

  private static final String URL_GET_EXECUTABLE_HISTORY = "/history?ajax=fetch";

  //executor apis
  private static final String URL_GET_EXECUTOR_STATS = "/serverStatistics";

  private String azBaseUrl;
  private String azuser;
  private String password;
  private volatile String sessionid;
  private volatile long logInTime;
  private long ttl = 5 * 60 * 1000;

  public AzkabanApi(String azBaseUrl, String azuser, String password) {
    this.azBaseUrl = azBaseUrl;
    this.azuser = azuser;
    this.password = password;
  }

  /**
   * login as azuser,return session id.
   *
   * @return session id
   */
  public String login() throws AzAuthException {
    String url = this.azBaseUrl + URL_LOGIN;
    Map<String, String> params = new HashMap<String, String>();
    params.put("username", this.azuser);
    params.put("password", this.password);
    try {
      Map<Object, Object> map = HttpCaller.post(Map.class, url, params, null, null);
      if (map.containsKey("error")) {
        throw new AzAuthException(map.get("error").toString());
      }
      sessionid = map.get("session.id").toString();
    } catch (IOException ioe) {
      throw new AzAuthException(ioe.getMessage());
    }
    logInTime = System.currentTimeMillis();
    return sessionid;
  }

  /**
   * create project on azkaban. when user create task,create project and upload task zip onto az.
   *
   * @param projectName project name
   * @param description project description
   */
  public void createProject(String projectName, String description) throws AzkabanException {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("name", projectName);
    params.put("description", description);
    params.put("session.id", this.sessionid);
    try {
      Map map = HttpCaller.post(Map.class, azBaseUrl + URL_CREATE_PROJECT, params, null, null);
      if (map.containsKey("error")) {
        throw new AzkabanException(map.get("message").toString(), null);
      }
    } catch (IOException ioe) {
      throw new AzkabanException("create project error", ioe);
    }

  }

  /**
   * where user create task,upload zip file onto az.
   *
   * @param projectZipFilePath zip file
   * @param project project name
   */
  public void uploadProject(String projectZipFilePath, String project) {
    checkLogin();
    RequestConfig.Builder requestConfigBuilder = RequestConfig.custom()
        .setConnectTimeout(10000)
        .setSocketTimeout(60 * 1000 * 5);
    CloseableHttpClient client = HttpClientBuilder.create()
        .setDefaultRequestConfig(requestConfigBuilder.build()).build();
    String url = this.azBaseUrl + URL_UPLOAD_PROJECT;
    HttpPost post = new HttpPost(url);
    try {
      File file = new File(projectZipFilePath);
      HttpEntity entity = MultipartEntityBuilder.create()
          .setMode(HttpMultipartMode.BROWSER_COMPATIBLE)
          .addBinaryBody("file", file, ContentType.create("application/zip"), file.getName())
          .addTextBody("project", project)
          .addTextBody("session.id", sessionid)
          .addTextBody("version", "1")
          .addTextBody("ajax", "upload")
          .build();
      post.setEntity(entity);
      CloseableHttpResponse response = client.execute(post);
      Header[] headers = post.getAllHeaders();
      int statusCode = response.getStatusLine().getStatusCode() / 100;
      String result = EntityUtils.toString(response.getEntity());
      response.close();
      if (statusCode == 4 || statusCode == 5) {
        throw new AzkabanException(result, null);
      }
      if (statusCode != 2) {
        throw new AzkabanException(result, null);
      }
    } catch (Exception e) {
      throw new AzkabanException("upload project failed", e);
    }
  }


  /**
   * remove project specified by name.
   *
   * @param name project name given by creation
   */
  public void deleteProject(String name) throws AzkabanException {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("project", name);
    params.put("session.id", this.sessionid);
    try {
      HttpCaller.get(String.class, azBaseUrl + URL_DELETE_PROJECT, params, null);
    } catch (IOException ioe) {
      throw new AzkabanException("delete project error", ioe);
    }
  }


  /**
   * fetch project flows of project.
   *
   * @param project project name
   * @return ProjectFlows
   */
  public ProjectFlows getProjectFlows(String project) {
    checkLogin();
    String url = azBaseUrl + URL_GET_PROJECT_FLOWS;
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("project", project);
    try {
      return HttpCaller.get(ProjectFlows.class, url, params, null);
    } catch (IOException ioe) {
      throw new AzkabanException("scheduler error", ioe);
    }
  }

  /**
   * get jobs of flow.
   *
   * @param project project name
   * @param flow flow id
   * @return FlowJobs
   */
  public FlowJobs getFlowJobs(String project, String flow) {
    checkLogin();
    String url = azBaseUrl + URL_GET_JOBS_OF_FLOW;
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("project", project);
    params.put("flow", flow);
    try {
      return HttpCaller.get(FlowJobs.class, url, params, null);
    } catch (IOException ioe) {
      throw new AzkabanException("scheduler error", ioe);
    }
  }


  /**
   * executions sorted in descendant submit time order.
   *
   * @param project project name
   * @param flow flow id
   * @param start start index
   * @param length max return count
   * @return FlowExecutions
   */
  public FlowExecutions getFlowExecutions(String project, String flow, int start, int length) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("flow", flow);
    params.put("project", project);
    params.put("start", start + "");
    params.put("length", length + "");
    try {
      return HttpCaller
          .get(FlowExecutions.class, azBaseUrl + URL_GET_FLOW_EXECUTIONS, params, null);
    } catch (Exception e) {
      throw new AzkabanException("call scheduler error", e);
    }

  }


  /**
   * getRunningExecutions.
   *
   * @param project project name
   * @param flow flow id
   * @return execids
   */
  public List<Integer> getRunningExecutions(String project, String flow) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("project", project);
    params.put("flow", flow);
    String url = azBaseUrl + URL_GET_FLOW_RUNNING_EXECUTIONS;
    try {
      Map<String, List<Integer>> map = HttpCaller.get(Map.class, url, params, null);
      return map.get("execIds");
    } catch (Exception e) {
      throw new AzkabanException("call scheduler error", e);
    }
  }


  /**
   * execute a project flow.
   *
   * @param project project name
   * @param flow flow id
   * @return FlowExecuteResponse
   */
  public FlowExecuteResponse executeFlow(String project, String flow,
      boolean bindExecutor,
      int executorId,
      FlowExecuteOptions options) {
    checkLogin();
    String url = this.azBaseUrl + URL_EXECUTE_FLOW;
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("project", project);
    params.put("flow", flow);
    if (options != null && !Strings.isNullOrEmpty(options.getConfigId())) {
      params.put("configId", options.getConfigId());
    }
    if (bindExecutor) {
      params.put("useExecutor", executorId + "");
    }
    try {
      return HttpCaller.post(FlowExecuteResponse.class, url, params, null, null);
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler failed", ioe);
    }
  }

  /**
   * cancels a running flow.if flow is not running,az will return error message, { "error" :
   * "Execution 302 of flow test isn't running." } but this method will ignore az error message.
   *
   * @param execid flow execution id
   */
  public void cancelFlowExecution(int execid) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("execid", execid + "");
    try {
      String url = azBaseUrl + URL_CALCEL_FLOW_EXECUTION;
      Map map = HttpCaller.get(Map.class, url, params, null);
      if (map != null && map.containsKey("error")) {
        throw new AzkabanException(map.get("error").toString(), null);
      }
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler error", ioe);
    }
  }

  /**
   * schedule a project flow.
   *
   * @param projectName project name
   * @param flowName flow name
   * @param cronExpression exec plan cron
   * @return scheduleid
   * @throws FlowScheduleException error returned by zakaban
   */
  public int scheduleCronFlow(String projectName, int projectId, String flowName,
      String cronExpression, boolean bindExecutor, int executorId, String concurrentOption) {
    checkLogin();
    String url = azBaseUrl + URL_SCHEDULE_FLOW;
    try {
      Map<String, String> params = new HashMap<String, String>();
      params.put("session.id", sessionid);
      params.put("flow", flowName);
      params.put("projectName", projectName);
      params.put("projectId", projectId + "");
      params.put("cronExpression", cronExpression);
      params.put("concurrentOption", concurrentOption);
      if (bindExecutor) {
        params.put("useExecutor", executorId + "");
      }
      Map<Object, Object> map = HttpCaller.post(Map.class, url, params, null, null);
      if (map.containsKey("error")) {
        throw new AzkabanException(map.get("error").toString(), null);
      }
      if (map.get("status").equals("error")) {
        throw new AzkabanException(map.get("message").toString(), null);
      } else {
        return Integer.parseInt(map.get("scheduleId").toString());
      }
    } catch (IOException e) {
      throw new AzkabanException("call scheduler error", e);
    }

  }

  /**
   * unschedule a flow.
   *
   * @param scheduleId schedule id
   */
  public void unscheduleFlow(int scheduleId) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("scheduleId", scheduleId + "");
    try {
      String url = azBaseUrl + URL_REMOVE_SCHEDULE;
      Map<String, String> map = HttpCaller.post(Map.class, url, params, null, null);
      if ("error".equals(map.get("status"))) {
        throw new AzkabanException(map.get("message").toLowerCase(), null);
      }
    } catch (IOException e) {
      throw new AzkabanException("call scheduler error", e);
    }

  }

  public FlowExecutionDetail getFlowExecution(int execid) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("execid", execid + "");
    try {
      String url = azBaseUrl + URL_GET_FLOW_EXEC_DETAIL;
      return HttpCaller.get(FlowExecutionDetail.class, url, params, null);
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler error", ioe);
    }
  }

  /**
   * fetch execution history.
   *
   * @param offset fetch offset
   * @param count max fetch count
   * @return history list
   */
  public List<AzFlowExecutionHistory> getFlowExecutionHistory(int offset, int count) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("offset", offset + "");
    params.put("count", count + "");
    try {
      String url = azBaseUrl + URL_GET_EXECUTABLE_HISTORY;
      String json = HttpCaller.get(String.class, url, params, null);
      Map map = JsonUtil.fromJson(Map.class, json);
      if (map.containsKey("error")) {
        throw new IOException(map.get("error").toString());
      }
      json = JsonUtil.toJson(map.get("history"));
      return JsonUtil.fromJsonList(AzFlowExecutionHistory.class, json);
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler error", ioe);
    }
  }

  /**
   * get execution log.
   *
   * @param execid The unique id for an execution
   * @param jobid The unique id for the job to be fetched
   * @param offset The offset for the log data
   * @param length The length of the log data
   * @return log
   */
  public JobLogData getExecutionJobLogs(int execid, String jobid, int offset, int length) {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("execid", String.valueOf(execid));
    params.put("jobId", jobid);
    params.put("offset", String.valueOf(offset));
    params.put("length", String.valueOf(length));
    try {
      String url = azBaseUrl + URL_GET_EXEC_JOB_LOGS;
      JobLogData data = HttpCaller.get(JobLogData.class, url, params, null);
      if (data.getLength() > 0) {
        data.setData(StringEscapeUtils.unescapeHtml4(data.getData()));
      }
      return data;
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler error", ioe);
    }
  }

  /**
   * get job params.
   *
   * @param project project name
   * @param flowid flow name
   * @param job job name
   * @return job params
   */
  public Map<String, String> getFlowJobParams(String project, String flowid, String job) {
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("project", project);
    params.put("flowName", flowid);
    params.put("jobName", job);
    try {
      String url = this.azBaseUrl + URL_GET_JOB_PARAMS;
      Map<String, Object> allParams = HttpCaller.get(Map.class, url, params, null);
      Map<String, String> result = new HashMap<>();
      for (Entry<String, Object> entry : allParams.entrySet()) {
        result.put(entry.getKey(), entry.getValue().toString());
      }
      result.remove("projectId");
      return result;
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler error", ioe);
    }
  }


  /**
   * set job params.
   *
   * @param project project name
   * @param flowid flow name
   * @param job job name
   * @param params param map
   */
  public void setFlowJobParams(String project, String flowid, String job,
      Map<String, String> params) {
    Map<String, String> finalParams = new HashMap<>();
    finalParams.put("session.id", sessionid);
    finalParams.put("project", project);
    finalParams.put("flowName", flowid);
    finalParams.put("jobName", job);
    params.remove("project");
    params.remove("jobName");
    params.remove("flowName");
    for (String key : params.keySet()) {
      finalParams.put("jobOverride[" + key + "]", params.get(key));
    }
    try {
      //jobOverride[execution-jar]
      Map<String, String> oldParams = this.getFlowJobParams(project, flowid, job);
      String url = this.azBaseUrl + URL_SET_JOB_PARAMS;
      Map<String, Object> map = HttpCaller.get(Map.class, url, finalParams, null);
      if (map != null && map.containsKey("error")) {
        throw new AzkabanException("scheduler error:" + map.get("error"), null);
      }
    } catch (IOException ioe) {
      throw new AzkabanException("call scheduler error", ioe);
    }
  }

  /**
   * get project flow schedule.
   *
   * @param project project name
   * @param flowid flowid
   * @return FlowSchedule
   */
  public FlowSchedule getFlowSchedule(String project, String flowid) {
    checkLogin();
    ProjectFlows flows = this.getProjectFlows(project);
    String projectId = flows.getProjectId();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    params.put("projectId", projectId + "");
    params.put("flowId", flowid);
    String url = azBaseUrl + URL_GET_SCHEDULE;
    try {
      Map<Object, Object> jsonMap = HttpCaller.get(Map.class, url, params, null);
      String json = JsonUtil.toJson(jsonMap.get("schedule"));
      if (json != null && json.length() > 0) {
        return JsonUtil.fromJson(FlowSchedule.class, json);
      } else {
        return null;
      }
    } catch (Exception e) {
      throw new AzkabanException("call scheduler error", e);
    }
  }

  /**
   * fetch all active azkaban executors.
   *
   * @return executors
   */
  public List<Executor> getAllActiveExecutors() {
    checkLogin();
    Map<String, String> params = new HashMap<String, String>();
    params.put("session.id", sessionid);
    String url = azBaseUrl + URL_GET_ACTIVE_EXECUTORS;
    try {
      Map<Object, Object> jsonMap = HttpCaller.get(Map.class, url, params, null);
      String json = JsonUtil.toJson(jsonMap.get("executors"));
      if (json != null && json.length() > 0) {
        List<Executor> executors = JsonUtil.fromJsonList(Executor.class, json);
        return executors;
      } else {
        return null;
      }

    } catch (Exception e) {
      throw new AzkabanException("call scheduler error", e);
    }
  }

  public ExecutorInfo getExecutorInfo(Executor executor) throws IOException {
    String url = "http://" + executor.getHost() + ":" + executor.getPort() + URL_GET_EXECUTOR_STATS;
    return HttpCaller.get(ExecutorInfo.class, url, null, null);
  }


  /**
   * get recently finished executions.
   *
   * @return execution list
   */
  public List<AzFlowExecutionHistory> getRecentlyFinished() {
    String url = azBaseUrl + URL_GET_RECENTLY_FINISHED;
    this.checkLogin();
    try {
      Map<String, String> params = new HashMap<String, String>();
      params.put("session.id", sessionid);
      Map<Object, Object> jsonMap = HttpCaller.get(Map.class, url, params, null);
      String json = JsonUtil.toJson(jsonMap.get("recentlyFinished"));
      if (json != null && json.length() > 0) {
        List<AzFlowExecutionHistory> executors = JsonUtil
            .fromJsonList(AzFlowExecutionHistory.class, json);
        return executors;
      } else {
        return Collections.emptyList();
      }

    } catch (Exception e) {
      throw new AzkabanException("call scheduler error", e);
    }

  }

  /**
   * URL_GET_RUNNING_FLOWS.
   *
   * @return runningflows
   */
  public List<AzFlowExecutionHistory> getRunningFlows() {
    String url = azBaseUrl + URL_GET_RUNNING_FLOWS;
    this.checkLogin();
    try {
      Map<String, String> params = new HashMap<String, String>();
      params.put("session.id", sessionid);
      Map<Object, Object> jsonMap = HttpCaller.get(Map.class, url, params, null);
      String json = JsonUtil.toJson(jsonMap.get("runningFlows"));
      if (json != null && json.length() > 0) {
        List<AzFlowExecutionHistory> executors = JsonUtil
            .fromJsonList(AzFlowExecutionHistory.class, json);
        return executors;
      } else {
        return Collections.emptyList();
      }

    } catch (Exception e) {
      throw new AzkabanException("call scheduler error", e);
    }
  }


  private void checkLogin() throws AzAuthException {
    if (System.currentTimeMillis() - this.logInTime > this.ttl) {
      this.login();
    }
  }
}
